import sys
import re
import os

from copy import copy
from functools import partial

import src.codeql_tools.viewer as codeql
import src.codeql_tools.preprocess as source
import src.codeql_tools.config as config



IsInit=False

SourceDir=""
SpecialName="def __init___0"



def extract_func_name(s):
    if s==SpecialName:
        return "__init___"
    s=s[:-1]
    s=s.split(' ')[1]
    return s

def init(index_path,codeql_query_dir,source_database_dir,test=False):
    global IsInit,SourceDir
    if not IsInit:
        if test:
            codeql.init_for_test(codeql_query_dir)
        else:
            codeql.init(codeql_query_dir)
        source.set_index_file(index_path)
        IsInit=True
        SourceDir=source_database_dir
        
def get_source_mapping():
    iterator=source.dataset_iter(SourceDir)
    for e in iterator:
        for index,error_number,source_path in e:
            mapping_data2ql={}
            mapping_data2source={}
            mapping_source_path=source_path[len(config.source_prefix):]
            mapping_source_path=f'{config.codeql_prefix}{mapping_source_path}'
            
            for function_name,(batch_number,data_path,data_number) in index:
                function_name=extract_func_name(function_name)
                mapping_data2ql[(data_path,data_number)]=codeql.Variable.construct_owner_id(mapping_source_path,['',function_name])
                mapping_data2source[(data_path,data_number)]=source_path,function_name
            yield mapping_data2ql,mapping_data2source

def get_codeql_validate_data(origin,sample_validate_data):
    def get_var(text:str):
        if text=="":
            return ""
        args=text.split('.')
        
        return args[0]
    
    error_msg=[]
    sample_vars,(error_key,occur_counter)=sample_validate_data
    validate_pass=True

        
    for e in sample_vars:
        if e not in origin:
            validate_pass=False
            error_msg.append(f'{len(error_msg)+1}. var: {e} miss.')
            
            
    if occur_counter>0:
        if error_key not in origin:
            validate_pass=False
            error_msg.append(f'{len(error_msg)+1}. error key: {error_key} miss.')
        elif origin[error_key]!=occur_counter:
            validate_pass=False
            error_msg.append(f'{len(error_msg)+1}. {error_key} counter error:  should be {occur_counter} but {origin[error_key]}.')
    return validate_pass,error_msg

def get_codeql_sample_data(origin_sample:source.Example,tokens,edges,counter):
    def new_error_location(origin,new_tokens,origin_group_by,new_group_by):
        if origin.has_bug:
            
            error_key=origin.tokens[origin.error_location]
            temp=origin_group_by[error_key]
            
            dev=len(new_group_by[error_key])-len(temp)
            
            return new_group_by[error_key][dev+temp.index(origin.error_location)]
        
        else:
            return 0
    def new_repair_candidate(origin,new_tokens,origin_group_by,new_group_by):
        vars={origin.tokens[x] for x in origin.repair_candidates if isinstance(x,int)}
        keywords=[x for x in origin.repair_candidates if isinstance(x,str)]
        result=[1]
        for v in vars:
            if v in new_group_by:
                result+=new_group_by[v]
            else:
                print(f"warn key:{v} not exists!")
        result+=keywords
        return result
    
    def new_repair_targets(origin,new_tokens,origin_group_by,new_group_by):
        vars={origin.tokens[x] for x in origin.repair_targets if isinstance(x,int)}
        keywords=[x for x in origin.repair_targets if isinstance(x,str)]
        result=[]
        for v in vars:
            if v in new_group_by:
                result+=new_group_by[v]
            else:
                result.append(1)
                print(f"warn key:{v} not exists!")
        result+=keywords
        return result
    
    validate_pass,error_logs=get_codeql_validate_data(counter,origin_sample.get_validate_info())
    if not validate_pass:
        return None
    
    
    new_tokens_group_by=source.Example.format_group_by_token(tokens)
    origin_group_by=origin_sample.get_group_by_info()
    
    # print(new_tokens_group_by)
    # print(origin_group_by)
    
    new_sample=copy(origin_sample)
    new_sample.edges=edges
    new_sample.error_location=new_error_location(origin_sample,tokens,origin_group_by,new_tokens_group_by)
    new_sample.repair_candidates=new_repair_candidate(origin_sample,tokens,origin_group_by,new_tokens_group_by)
    new_sample.repair_targets=new_repair_targets(origin_sample,tokens,origin_group_by,new_tokens_group_by)
    new_sample.tokens=tokens
    return new_sample

def validate_index():
    logger_file=open('running_detail.log','w')
    
    #init(config.index_path,config.query_dir,config.code_source,False)
    
    source.validate(config.code_source)
    
    
    no_exists_counter=0
    miss_function_name=[]
    total=0
    source_path=""
    for a,b in get_source_mapping():  
        round_no_exists_counter=0
        keys=list(a.keys())[:]
        for key in keys:
            total+=1
            data_path,data_number=key
            source_path,function_name=b[key]
            if a[key] in codeql.view.FuncViewByOwnerId:
                pass
            else:
                no_exists_counter+=1
                miss_function_name.append(function_name)
                round_no_exists_counter+=1
                print(f'{function_name}|{a[key]} in {source_path} miss')  
            # if round_no_exists_counter>0:
            #     vkeys=codeql.FuncViewByOwnerId.keys()
                
            #     for name in miss_function_name:
            #         print(f'{name} in {source_path} miss,may be in:')   
            #         for key in vkeys:         
            #             if key.find(name)!=-1:
            #                 print(key)
        print(f"in {source_path} index failure: {len(keys)}/{round_no_exists_counter}")
    if no_exists_counter>0:
        print("index validate not pass.")
        example_keys=list(codeql.view.FuncViewByOwnerId.keys())[:100]
        for e in example_keys:
            print(e)
    print(f"total index success: {total}/{total-no_exists_counter}")
    logger_file.write(f"toal index success: {total}/{total-no_exists_counter}+\n")
    
    
    
    logger_file.close()

def validate_token():
    logger_file=open('running_detail.log','w')
    
    #init(config.index_path,config.query_dir,config.code_source)
    validate_pass_counter=0
    total=0
    source_path=""
    for data2ql,data2source in get_source_mapping():  
        round_validate_pass_counter=0
        keys=list(data2ql.keys())[:]
        for key in keys:
            total+=1
            data_path,data_number=key
            source_path,function_name=data2source[key]
            if data2ql[key] in codeql.view.FuncViewByOwnerId:
                func_obj=codeql.view.FuncViewByOwnerId[data2ql[key]]
                tokens,counter=codeql.view.tokeize(func_obj,False,False)
                example=source.get_examples(data_path)[data_number]
                example=source.Example(example)
                
                validate_pass,error_logs=get_codeql_validate_data(counter,example.get_validate_info())
                if validate_pass:
                    validate_pass_counter+=1
                    round_validate_pass_counter+=1
                else:
                    logger_file.write(f'in dataset file {data_path}:{data_number}+\n')
                    print(f'in dataset file {data_path}:{data_number}')
                    logger_file.write(f'in source file {data2source[key]}+\n')
                    logger_file.write(f"{tokens}+\n")
                    logger_file.write(f"{example.get_repair_info()}")
                    print(f'in source file {data2source[key]}')
                    for log in error_logs:
                        logger_file.write(log+'\n')
                        print(log)
        print(f'in {source_path},validate success: {len(keys)}/{round_validate_pass_counter}')  
            # if round_no_exists_counter>0:
            #     vkeys=codeql.FuncViewByOwnerId.keys()
                
            #     for name in miss_function_name:
            #         print(f'{name} in {source_path} miss,may be in:')   
            #         for key in vkeys:         
            #             if key.find(name)!=-1:
            #                 print(key)
    print(f"validate success: {total}/{validate_pass_counter}")
    logger_file.write(f"validate success: {total}/{validate_pass_counter}+\n")
    
    logger_file.close()

def generate_sample():
    generate_success_counter=0
    total=0
    for data2ql,data2source in get_source_mapping():  
        round_validate_pass_counter=0
        keys=list(data2ql.keys())[:]
        for key in keys:
            total+=1
            data_path,data_number=key
            source_path,function_name=data2source[key]
            if data2ql[key] in codeql.view.FuncViewByOwnerId:
                func_obj=codeql.view.FuncViewByOwnerId[data2ql[key]]
                tokens,edges,counter=codeql.view.tokens_and_cfg(func_obj)
                example=source.get_examples(data_path)[data_number]
                example=source.Example(example)
                
                new_example=get_codeql_sample_data(example,tokens,edges,counter)
                if new_example:
                    generate_success_counter+=1
                    round_validate_pass_counter+=1
                    yield new_example
        print(f'in {source_path},generate sample success: {len(keys)}/{round_validate_pass_counter}')  
    print(f"total sample success: {total}/{generate_success_counter}")

def extract_codeql_sample(dir,baseFilename,batch_number=5000):
    def get_batch_fp(now_batch,now_number,now_fp,now_filename):
        if now_filename is None:
            new_filename=f'{dir}/{baseFilename}_{now_batch}'
            print(f"new batch: {new_filename}")
            return open(new_filename,'w'),new_filename,0,0
        
        elif now_number==batch_number:
            now_fp.close()
            new_filename=f'{dir}/{baseFilename}_{now_batch+1}'
            print(f"new batch: {new_filename}")
            return open(new_filename,'w'),new_filename,now_batch+1,0
        
        else:
            return now_fp,now_filename,now_batch,now_number+1
        
    filename=None    
    fp=None
    now_batch=0
    now_number=0
    
    if not os.path.exists(dir):
        os.makedirs(dir)
    
    for sample in generate_sample():
           fp,filename,now_batch,now_number=get_batch_fp(now_batch,now_number,fp,filename) 
           fp.write(f'{sample.to_json()}\n')
    
    if fp:
        fp.close()              
    

if __name__=="__main__":
    args_length=len(sys.argv)
    if args_length==1:
        extract_codeql_sample(config.output_dir,config.output_baseFileName,config.output_batch)
    elif args_length==2:
        if sys.argv[1]=="test_index":
            validate_index()
        elif sys.argv[1]=="test_token":
            validate_token()
        else:
            print("input error.")
    